import torch
from torch import nn

def weights_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
        # m.weight.data.normal_(0.0, 0.02)
        nn.init.kaiming_normal_(m.weight)
        m.bias.data.fill_(0)

class Generator(nn.Module):
    #生成器
    def __init__(self, input_size):
        super(Generator, self).__init__()

        self.leakrelu = nn.LeakyReLU()
        # [b, 100] => [b, 3*3*512]
        self.fc = nn.Linear(input_size, 3*3*512)

        self.conv1 = nn.ConvTranspose2d(512, 256, 3, 3)
        self.bn1 = nn.BatchNorm2d(256)


        self.conv2 = nn.ConvTranspose2d(256, 128, 5, 2)
        self.bn2 = nn.BatchNorm2d(128)

        self.conv3 = nn.ConvTranspose2d(128, 3, 4, 3)
        self.tanh = nn.Tanh()
    def forward(self, inputs):

        x = self.fc(inputs)
        # [b, 512, 3, 3]
        x = x.view(-1, 512, 3, 3)
        x = self.leakrelu(x)
        # [b, 256, 9, 9]
        x = self.leakrelu(self.bn1(self.conv1(x)))
        # [b, 128, 21, 21]
        x = self.leakrelu(self.bn2(self.conv2(x)))
        # [b, 3, 64, 64]
        x = self.tanh(self.conv3(x))

        return x

class Discriminator(nn.Module):
    # 判别器
    def __init__(self, input_channel):
        super(Discriminator, self).__init__()

        self.leakrelu = nn.LeakyReLU()

        self.conv1 = nn.Conv2d(input_channel, 64, 5, 3)

        self.conv2 = nn.Conv2d(64, 128, 5, 3)
        self.bn2 = nn.BatchNorm2d(128)

        self.conv3 = nn.Conv2d(128, 256, 5, 3)
        self.bn3 = nn.BatchNorm2d(256)

        self.fc = nn.Linear(256, 1)


    def forward(self, inputs):
        # [b, 64, 20, 20]
        x = self.leakrelu(self.conv1(inputs))
        # [b, 128, 6, 6]
        x = self.leakrelu(self.bn2(self.conv2(x)))
        # [b, 256, 1, 1]
        x = self.leakrelu(self.bn3(self.conv3(x)))
        # [b, 256]
        x = x.view(inputs.shape[0], -1)
        # [b, 1]
        x = self.fc(x).squeeze()

        return x

def main():

    data = torch.rand(128, 100)
    genertor = Generator(100)
    genertor.apply(weights_init)
    out = genertor(data)
    print(out.shape)

    discriminator = Discriminator(3)
    discriminator.apply(weights_init)
    out = discriminator(out)
    print(out.shape)

if __name__ == '__main__':
    main()